import openai
import myopenaiapikey
import json
openai.api_base = myopenaiapikey.base
openai.api_key = myopenaiapikey.api



def usemsf(ip:str,method:str):
    print("正在使用msf：参数为："+ip+","+method)

def usenmap(ip:str):
    print("正在使用nmap：参数为：" + ip)

def usesqlmap(url:str):
    print("正在使用sqlmap：参数为：" + str(url))





functions2 = [
    {
    "name":"usemsf",
    "description":"""
    返回选择的渗透测试工具为metasploit-framework,简称msf
    ip这个参数为想要渗透的目标的ip地址
    method这个参数有三个可选项
    如果:“method”为“expolit”则表示选择漏洞利用模块对目标进行渗透
    如果:“method”为“auxliary”则表示选择漏洞扫描模块对目标进行扫描
    如果:“method”为“handler”则表示选择生成木马对目标钓鱼攻击
    """,
    "parameters":{
        "type":"object",
        "properties": {
                "ip":{
                        "type":"string",
                        "description":"想要渗透的目标的ip地址"
                    },
                "method":
                    {
                        "type":"string",
                        "description":"想要渗透的方式"
                    },
            },
        
        
        "required":["ip","method"]
    }

},
{
    "name":"usenmap",
    "description":"""
    返回选择的渗透测试工具为nmap,是一款端口扫描工具
    ip这个参数为想要渗透的目标的ip地址
    """,
    "parameters":{
        "type":"object",
        "properties": {
        "ip":{
            "type":"string",
            "description":"想要渗透的目标的ip地址"
            },
        },
        "required":["ip"]
    }

},
{
    "name":"usesqlmap",
    "description":"""
    返回选择的渗透测试工具为sqlmap,是一款web应用扫描工具
    url这个参数为想要渗透的目标的url地址
    """,
    "parameters":{
        "type":"object",
        "properties": {
        "url":{
            "type":"string",
            "description":"想要渗透的目标的url地址"
            },
        },
        "required":["url"]
    }

}


]

def chat2(query):
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo-0613",
        messages=[{"role": "user", "content": query}],
        functions=functions2,
    )
    message = response["choices"][0]["message"]
    return message



query = "我想对目标http://192.168.2.5做sql注入"
message = chat2(query)
print(message)
function_name = message["function_call"]["name"]
print(function_name)

if function_name == "usemsf":
    ip = json.loads(message["function_call"] ["arguments"]).get("ip")
    method = json.loads(message["function_call"] ["arguments"]).get("method")
    print(ip+"\r\n"+method)
    usemsf(ip=ip,method=method)
elif function_name == "usenmap":
    ip = json.loads(message["function_call"] ["arguments"]).get("ip")
    print(ip)
    usenmap(ip=ip)
elif function_name == "usesqlmap":
    url = json.loads(message["function_call"] ["arguments"]).get("url")
    print(url)
    usesqlmap(url=url)


